from mindspore import nn, ops

"""Distilling the Knowledge in a Neural Network"""
class KLLoss(nn.LossBase):
    def __init__(self, T, reduction="sum"):
        super(KLLoss, self).__init__(reduction)
        self.kldiv_loss = ops.KLDivLoss(reduction="none")
        self.log_softmax = ops.LogSoftmax(axis=1)
        self.softmax = ops.Softmax(axis=1)
        self.T = T

    def construct(self, y_s, y_t):
        p_t = self.log_softmax(y_t / self.T)
        p_s = self.softmax(y_s / self.T)
        x = self.kldiv_loss(p_t, p_s)
        x = self.get_loss(x) / y_s.shape[0]
        return x


class KLLossCell(nn.Cell):
    def __init__(self, student_net, T):
        super(KLLossCell, self).__init__(auto_prefix=False)
        self.student_net = student_net
        self.kl_loss = KLLoss(T)
        self.ce_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')

    def construct(self, teacher_fm, teacher_logits, data, label):
        student_logits = self.student_net(data)
        return self.ce_loss(student_logits, label) \
            + self.kl_loss(student_logits, teacher_logits)
    
    @property
    def backbone_network(self):
        return self.student_net
